# main.py
import torch
import torch.nn.functional as F
from dataLoader import load_datasets
from model import NetMNIST, NetCIFAR
from GPDC import gpdc

def main():
    # 加载数据集
    train_loader_mnist, test_loader_mnist = load_datasets('mnist')
    train_loader_cifar, test_loader_cifar = load_datasets('cifar')

    # 初始化模型
    model_mnist = NetMNIST()
    model_cifar = NetCIFAR()

    # 训练和评估MNIST模型
    print("Training MNIST Model...")
    gpdc(model_mnist, train_loader_mnist, epochs=10)
    evaluate(model_mnist, test_loader_mnist)

    # 训练和评估CIFAR模型
    print("Training CIFAR Model...")
    gpdc(model_cifar, train_loader_cifar, epochs=10)
    evaluate(model_cifar, test_loader_cifar)

def evaluate(model, test_loader):
    correct = 0
    total = 0
    model.eval()
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network on the test images: {100 * correct / total}%')

if __name__ == "__main__":
    main()